from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from database import Base, DATABASE_URL
import os
import shutil

def clean_database():
    # 创建数据库引擎
    engine = create_engine(DATABASE_URL)
    
    # 创建会话
    Session = sessionmaker(bind=engine)
    session = Session()
    
    try:
        # 删除所有表
        Base.metadata.drop_all(engine)
        print("所有表已删除")
        
        # 重新创建所有表
        Base.metadata.create_all(engine)
        print("所有表已重新创建")
        
        # 清理上传文件夹
        uploads_dir = "uploads"
        if os.path.exists(uploads_dir):
            shutil.rmtree(uploads_dir)
            os.makedirs(uploads_dir)
            print("上传文件夹已清理")
        
        print("数据库清理完成！")
        
    except Exception as e:
        print(f"清理过程中出现错误: {str(e)}")
        session.rollback()
    finally:
        session.close()

if __name__ == "__main__":
    # 确认操作
    confirm = input("此操作将删除所有数据，是否继续？(y/n): ")
    if confirm.lower() == 'y':
        clean_database()
    else:
        print("操作已取消") 